# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime
from enum import IntEnum
import json
import os

import pandas as pd
from pandas.io.sql import to_sql
import jinja2
import MySQLdb

from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.strategy.mm.subscription import (SubscriptionRequest, get_subscription_key)
from coin.support.feed_tool.feed_stats.logic.util import (
    get_exchange_api_id_from_request,
    gen_engine,
    gen_engine_for_importer,
    query_quote_to_usd_exchange_rate,
)

THRESHOLD_MAP = {
    'no_book': [
        ('no_book_duration', 3000),  # second
        ('volume_usd', 100000),
        ('book_count', 20000)
    ],
    'no_hour_book': [
        ('no_hourly_book_count', 0),  # all hours must not be empty
        ('book_count', 200)
    ],
    'invalid_book_ratio': [
        ('book_invalid_count/book_count', 0.01),  # no more than 1% invalid book
        ('volume_usd', 100000),
        ('book_count', 20000)
    ],
    'no_trade': [
        ('no_trade_duration', 5000),  # second
        ('volume_usd', 100000),
        ('trade_count', 3000)
    ],
    'between_book_trade': [('max/avg_between_book_trade_count', 60),
                           ('max_between_book_trade_duration', 5)],
    'outside_bid_ask_trade': [('max/avg_outside_bid_ask_trade_count', 60),
                              ('max_between_book_trade_count', 2000)],
    'trade_book_ratio':
        20,
    'feed_checker_consecutive_invalid': [
        ('book_max_feed_checker_consecutive_invalid_count', 10),
        ('book_max_feed_checker_consecutive_invalid_period', 1800),  # second
        ('volume_usd', 100000),
        ('book_count', 20000)
    ],
}

QUOTE_LIST = [
    'USD',
    'EUR',
    'KRW',
    'JPY',
    'USDT',
    'BTC',
    'ETH',
]


class FeedHealthResult(IntEnum):
  GOOD = 1
  BAD = 2
  FEED_NOT_EXIST = 3


class FeedStatsNotReadyError(RuntimeError):
  pass


def _gen_query_feed_stats_to_check_health_sql(trading_date, current_date, git_commit_sha):
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  machines = [
      "feed-01.ap-northeast-2.aws",
      "feed-04.ap-northeast-1.aws",
  ]
  sql_tmpl = 'coin/support/feed_tool/feed_stats/query_sql/' \
             'query_feed_stats_to_check_health.sql.tmpl'
  template = template_env.get_template(sql_tmpl)
  current_date = current_date or datetime.datetime.now().strftime('%Y-%m-%d')
  return template.render(trading_date=trading_date,
                         git_commit_sha=git_commit_sha,
                         current_date=current_date,
                         machines=machines)


def query_feed_stats_to_check_health(trading_date, current_date, git_commit_sha):
  engine = gen_engine()
  feed_stats = pd.read_sql(
      _gen_query_feed_stats_to_check_health_sql(trading_date, current_date, git_commit_sha), engine)
  return feed_stats


def filter_and_gen_columns(feed_stats):
  feed_stats['max_between_book_trade_duration'] = \
      feed_stats['max_between_book_trade_duration'].fillna(0)
  feed_stats['max_outside_bid_ask_trade_duration'] = \
    feed_stats['max_outside_bid_ask_trade_duration'].fillna(0)
  feed_stats['book_max_feed_checker_consecutive_invalid_count'] = \
    feed_stats['book_max_feed_checker_consecutive_invalid_count'].fillna(0)
  feed_stats['max_outside_bid_ask_trade_duration'] = \
    feed_stats['max_outside_bid_ask_trade_duration'].fillna(0)
  feed_stats.dropna(how='any', inplace=True)
  feed_stats = feed_stats.sort_values('create_time').groupby(
      ['machine', 'exchange', 'trading_date', 'symbol'],
      as_index=False).last().reset_index(drop=True)
  feed_stats = feed_stats.loc[feed_stats['quote'].isin(QUOTE_LIST)]
  pd.options.mode.chained_assignment = None  # default='warn'
  exchange_rate = query_quote_to_usd_exchange_rate(QUOTE_LIST)
  feed_stats['quote_exchange'] = feed_stats['quote'].map(exchange_rate)
  feed_stats['volume_usd'] = feed_stats['volume'] * feed_stats['quote_exchange']
  # Check hourly book count, if there are 0 case, feed is wrong.
  # For Upbit, there is a special feed error case that there are exactly 142
  # books in an hour. Mark this case as no_hourly_book_count as well.
  feed_stats['no_hourly_book_count'] = feed_stats['book_hour_count'].str.count(' 0,') + \
      feed_stats['book_hour_count'].str.count(' 142,')
  feed_stats['book_invalid_count/book_count'] = (feed_stats['book_feed_checker_invalid_count']
                                                 / feed_stats['book_count'])
  feed_stats['max/avg_between_book_trade_count'] = (feed_stats['max_between_book_trade_count']
                                                    / feed_stats['avg_between_book_trade_count'])
  feed_stats['max/avg_outside_bid_ask_trade_count'] = (
      feed_stats['max_outside_bid_ask_trade_count'] / feed_stats['avg_outside_bid_ask_trade_count'])
  return feed_stats


def find_unhealthy_feed_stats(feed_stats, threshold):
  threshold_name, threshold_value = threshold
  if not isinstance(threshold_value, list):
    return feed_stats.loc[feed_stats[threshold_name] > threshold_value]
  else:
    for _threshold_name, _threshold_value in threshold_value:
      feed_stats = feed_stats.loc[feed_stats[_threshold_name] > _threshold_value]
    return feed_stats


def _gen_query_feed_health_stats_sql(trading_date):
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  sql_tmpl = 'coin/support/feed_tool/feed_stats/query_sql/' \
             'query_feed_health_stats_for_1day.sql.tmpl'
  template = template_env.get_template(sql_tmpl)
  return template.render(trading_date=trading_date)


def query_old_feed_health_stats(trading_date):
  engine = gen_engine()
  feed_stats = pd.read_sql(_gen_query_feed_health_stats_sql(trading_date), engine)
  return feed_stats


def gen_latest_feed_health_stats_with_raw(latest_feed_health_stats, raw_feed_stats):
  feed_not_exist_raw_stats = raw_feed_stats.loc[raw_feed_stats['book_count'] == 0]
  feed_not_exist_raw_stats['health'] = FeedHealthResult.FEED_NOT_EXIST
  feed_not_exist_raw_stats['reason'] = None
  picked_columns = ['trading_date', 'machine', 'exchange_api_id', 'symbol', 'health', 'reason']
  feed_not_exist_raw_stats = feed_not_exist_raw_stats[picked_columns]
  latest_feed_health_stats = pd.concat([latest_feed_health_stats, feed_not_exist_raw_stats])
  return latest_feed_health_stats


def _gen_delete_feed_health_stats_sql(trading_date):
  template_loader = jinja2.FileSystemLoader(searchpath="./")
  template_env = jinja2.Environment(loader=template_loader)

  sql_tmpl = 'coin/support/feed_tool/feed_stats/query_sql/' \
             'delete_feed_health_stats.sql.tmpl'
  template = template_env.get_template(sql_tmpl)
  return template.render(trading_date=trading_date)


def _insert_feed_health_db(feed_health_stats):
  engine = gen_engine_for_importer()
  to_sql(feed_health_stats, 'FeedHealthStats', engine, if_exists='append', index=False)


def insert_unhealthy_to_feed_health_stats_db(trading_date, latest_feed_health_stats,
                                             db_config_file):
  db_config_path = os.path.expanduser(db_config_file)
  assert db_config_path and os.path.isfile(db_config_path), \
      'DB config file does not exist: %s' % db_config_path

  with open(db_config_path, 'r') as open_file:
    db_config = json.load(open_file)
  conn = MySQLdb.connect(host=db_config.get('host', None),
                         db=db_config.get('database', None),
                         user=db_config.get('user', None),
                         passwd=db_config.get('password', None))
  cursor = conn.cursor()

  sql = _gen_delete_feed_health_stats_sql(trading_date)
  cursor.execute(sql)
  conn.commit()
  conn.close()

  _insert_feed_health_db(latest_feed_health_stats)


def update_feed_health_stats(latest_unhealthy_stats, raw_feed_stats, trading_date, db_config_file):
  latest_feed_health_stats = gen_latest_feed_health_stats_with_raw(latest_unhealthy_stats,
                                                                   raw_feed_stats)
  insert_unhealthy_to_feed_health_stats_db(trading_date, latest_feed_health_stats, db_config_file)


def print_feed_health_check_result(unhealthy_stats, threshold):
  if len(unhealthy_stats) == 0:
    return

  threshold_name, threshold_value = threshold
  print('Unhealthy feed stats %s > %s:' % (threshold_name, threshold_value))
  print(unhealthy_stats.to_string(float_format=lambda x: '%.2f' % x, index=False))
  print()


def integrate_unhealthy_stats(latest_feed_health_stats, unhealthy_stats, threshold):
  picked_columns = ['trading_date', 'machine', 'exchange_api_id', 'symbol']
  unhealthy_stats = unhealthy_stats[picked_columns]
  unhealthy_stats['health'] = FeedHealthResult.BAD
  unhealthy_stats['reason'] = threshold[0]  # threshold name
  unhealthy_stats.reset_index(drop=True, inplace=True)

  latest_feed_health_stats = pd.concat([latest_feed_health_stats, unhealthy_stats],
                                       axis=0,
                                       ignore_index=True)
  picked_columns += ['health']
  latest_feed_health_stats = latest_feed_health_stats.groupby(picked_columns, as_index=False).\
      agg({'reason': lambda x: ','.join(x)}).reset_index(drop=True)
  return latest_feed_health_stats


def run_feed_health_check(trading_date, current_date, git_commit_sha, db_config_file):
  raw_feed_stats = query_feed_stats_to_check_health(trading_date, current_date, git_commit_sha)

  if len(raw_feed_stats.loc[raw_feed_stats['book_count'] > 0]) < 300:
    latest_feed_health_stats = pd.DataFrame()
    update_feed_health_stats(latest_feed_health_stats, raw_feed_stats, trading_date, db_config_file)
    raise FeedStatsNotReadyError('Feed not ready for %s on %s.' % (git_commit_sha, trading_date))

  feed_stats = filter_and_gen_columns(raw_feed_stats.copy())
  latest_feed_health_stats = pd.DataFrame()
  for threshold in THRESHOLD_MAP.items():
    unhealthy_stats = find_unhealthy_feed_stats(feed_stats, threshold)
    print_feed_health_check_result(unhealthy_stats, threshold)
    latest_feed_health_stats = integrate_unhealthy_stats(latest_feed_health_stats,
                                                         unhealthy_stats,
                                                         threshold)

  update_feed_health_stats(latest_feed_health_stats, raw_feed_stats, trading_date, db_config_file)


def query_feed_health(trading_dates, machine, market_type, exchange, api_version, symbol):
  sql_tmpl = """
SELECT
*
FROM feed_stat_20200803.FeedHealthStats
WHERE trading_date in ({trading_dates})
AND machine = '{machine}'
AND exchange_api_id = {exchange_api_id}
AND symbol = '{symbol}';
"""

  request_list = [SubscriptionRequest(market_type, exchange, api_version)]
  exchange_api_id = get_exchange_api_id_from_request(request_list)[0]
  trading_date_strs = [trading_date.strftime('%Y%m%d') for trading_date in trading_dates]
  sql = sql_tmpl.format(trading_dates=','.join(trading_date_strs),
                        machine=machine,
                        exchange_api_id=exchange_api_id,
                        symbol=symbol)

  engine = gen_engine()
  feed_stats = pd.read_sql(sql, engine)
  return feed_stats


def find_feed_healthy_trading_dates(trading_dates,
                                    machine,
                                    subscription_request: SubscriptionRequest,
                                    symbol):
  feed_healthy_trading_dates = []
  feed_health_stats = query_feed_health(trading_dates,
                                        machine,
                                        subscription_request.market_type,
                                        subscription_request.exchange,
                                        subscription_request.api_version,
                                        symbol)
  for trading_date in trading_dates:
    feed_health_stats_1day = feed_health_stats.loc[feed_health_stats['trading_date'] ==
                                                   trading_date.date()]
    if len(feed_health_stats_1day) == 0:
      feed_healthy_trading_dates.append(trading_date)
    elif len(feed_health_stats_1day) == 1:
      if feed_health_stats_1day.iloc[0]['health'] == FeedHealthResult.GOOD:
        feed_healthy_trading_dates.append(trading_date)
      elif (feed_health_stats_1day.iloc[0]['reason'] == 'no_hour_book'
            and (subscription_request.exchange == 'Okex'
                 and subscription_request.market_type == 'Futures')
            and trading_date.weekday() == 4):
        # Exception for Okex futures Friday rollover case.
        feed_healthy_trading_dates.append(trading_date)
    else:
      raise ValueError(feed_health_stats_1day)
  return feed_healthy_trading_dates


def filter_common_good_trading_dates(trading_dates, machine, products):
  common_good_trading_dates = set(trading_dates)
  for product in products:
    good_trading_dates = []
    subscription_request = get_subscription_key(product)
    for trading_date in trading_dates:
      if (subscription_request.exchange == 'Okex'
          and subscription_request.market_type == 'Futures'):
        product = OkexFuturesProduct.FromCurrencyPair2AndContractType(product.base,
                                                                      product.quote,
                                                                      product.contract_type,
                                                                      trading_date)
      good_trading_dates += find_feed_healthy_trading_dates([trading_date],
                                                            machine,
                                                            subscription_request,
                                                            product.symbol)
    common_good_trading_dates.intersection_update(set(good_trading_dates))
  return sorted(list(common_good_trading_dates))
